import cPickle as pickle
from astro.io import readtxt
from astro.coord import ang_sep
import numpy as np
import matplotlib.pyplot as pl
from matplotlib.mlab import rec_append_fields

import xcorr
from xcorr import \
     wlya, wlyb, Ckms, cosmo, PC, hires_filenames

wc4 = 1548.2041

# maximum separation allowed for an LBG from a QSO sight-line
# (degrees).
MAXSEP = 40. / 60.

# bins to use when measuring mean flux. distances are in Mpc
maxdist = 50.0

# linear bins
#BINEDGES = np.arange(0, maxdist + binwidth, binwidth)

# log bins
BINEDGES = np.logspace(0, 3, 31)

NBINS = len(BINEDGES) - 1
#NBINS_R = len(radial_binedges) - 1
#NBINS_T = len(transverse_binedges) - 1
#CBINS = BINEDGES[:-1] + BINEDGES[1:] - BINEDGES[:-1]

CBINS = BINEDGES[:-1] + 0.5*(BINEDGES[1:] - BINEDGES[:-1])


# plotting
A4LANDSCAPE = 11.7, 8.3

NRAND = 1000


def calc_npairs_t(bins,
                 lbgs_angdist, lbgs_losdist, lbgs_num,
                 abslosdist, absnums):
    """ Find the number of abs-LBG pairs as a function of their
    transverse separation (only include pairs that are within some
    arbitrary maximum radial separation).
    """
    nbins = len(bins) - 1

    # number of galaxy-abs pairs per bin
    
    # pair counts
    pc = np.zeros(nbins, float)
    idpairs = [[] for i in range(nbins)]
    # for each abs system, calculate (1) lod and transverse separation
    # for each nearby lbg and (2) add pairs to the appropriate pc bins

    for abslos, absnum in zip(abslosdist, absnums):
        # note we have calculated the transverse comoving distance
        # (angdist) of an LBG from the QSO sightline at the redshift
        # of the LBG. This will introduce an error that increases with
        # the LOS comoving distance difference between the absorber
        # and the galaxy.

        sep_r = np.abs(abslos - lbgs_losdist)
        sep_t = lbgs_angdist

        # only keep lbgs that are within arbitrary maximum radial
        # distance separation:
        isclose = sep_r < 15.0 

        sep_t = sep_t[isclose]
        lbgnum = lbgs_num[isclose]

        # this loop may be able to be replaced by numpy index trickery
        # for increased speed. It's ok for now though.
        for t,num in zip(sep_t, lbgnum):
            # find the r, t bin in which this pair lies
            ind = bins.searchsorted(t)
            #print r,t, ir, it, rbins[ir-1:ir+1], tbins[it-1:it+1]
            if ind > nbins:
                # outside largest bin
                continue
            # increment pairs in bin by one
            ind -= 1
            pc[ind] += 1
            idpairs[ind].append( (num, absnum) )

    return pc, idpairs
    

def calc_npairs_s(bins,
                 lbgs_angdist, lbgs_losdist, lbgs_num,
                 abslosdist, absnums, maxdist):
    """ Find the number of abs-LBG pairs as a function of their 3-d
    separation.
    """
    nbins = len(bins) - 1

    # number of galaxy-abs pairs per bin
    
    # pair counts
    pc = np.zeros(nbins, float)
    idpairs = [[] for i in range(nbins)]
    # for each abs system, calculate (1) lod and transverse separation
    # for each nearby lbg and (2) add pairs to the appropriate pc bins

    for abslos, absnum in zip(abslosdist, absnums):
        # note we have calculated the transverse comoving distance
        # (angdist) of an LBG from the QSO sightline at the redshift
        # of the LBG. This will introduce an error in the 3d distance
        # that increases with the LOS comoving distance difference
        # between the absorber and galaxy.
        
        sep = np.hypot(lbgs_angdist, abslos - lbgs_losdist)
        # only keep lbgs that are within arbitrary maximum radial
        # distance separation:
        isclose = sep < maxdist

        lbgnum = lbgs_num[isclose]
        sep = sep[isclose]

        # this loop may be able to be replaced by numpy index trickery
        # for increased speed. It's ok for now though.
        for s,num in zip(sep, lbgnum):
            # find the s bin in which this pair lies
            ind = bins.searchsorted(s)
            if ind > nbins:
                # outside largest bin
                continue
            # increment pairs in bin by one
            ind -= 1
            pc[ind] += 1
            idpairs[ind].append( (num, absnum) )

    return pc, idpairs
    
def calc_npairs(rbins, tbins,
               lbgs_angdist, lbgs_losdist, lbgs_num,
               abslosdist, absnums):
    """ Find the number of abs-LBG pairs as a function of their radial
    and transverse separations. Generates a grid of pairs.
    """
    nbins_r = len(rbins) - 1
    nbins_t = len(tbins) - 1

    # number of galaxy-abs pairs per bin
    
    # pair counts
    pc = np.zeros([nbins_r, nbins_t], float)
    idpairs = list_matrix(nbins_r, nbins_t)
    # for each abs system, calculate (1) lod and transverse separation
    # for each nearby lbg and (2) add pairs to the appropriate pc bins

    for abslos, absnum in zip(abslosdist, absnums):
        # note we have calculated the transverse comoving distance
        # (angdist) of an LBG from the QSO sightline at the redshift
        # of the LBG. This will introduce an error that increases with
        # the LOS comoving distance difference between the absorber
        # and the galaxy.

        sep_r = np.abs(abslos - lbgs_losdist)
        sep_t = lbgs_angdist

        # this loop may be able to be replaced by numpy index trickery
        # for increased speed. Ok for now though
        for r,t,num in zip(sep_r, sep_t, lbgs_num):
            # find the r, t bin in which this pair lies
            ir = rbins.searchsorted(r)
            it = tbins.searchsorted(t)
            #print r,t, ir, it, rbins[ir-1:ir+1], tbins[it-1:it+1]
            if ir > nbins_r or it > nbins_t:
                # outside largest bin
                continue
            # increment pairs in bin by one
            i,j = ir - 1, it - 1
            pc[i, j] += 1
            idpairs[i][j].append( (num, absnum) )

    return pc, idpairs

def list_matrix(n,m):
    return [ [[] for i in range(n)] for j in range(m)]

def list_matrix_extend(mat1, mat2):
    len1 = len(mat1)
    len2 = len(mat1[0])
    for i in range(len1):
        for j in range(len2):
            mat1[i][j].extend(mat2[i][j])


#def xcorr_setup():
#    pass

def readCIV():
    """ Read all the CIV systems."""
    CIVa = readtxt('/home/nhmc/projects/gas_gal_hiz/CIV/CIV_all.txt',
                  readnames=True)

    const = 1. / (PC * 1.e6)

    # 1. get the detection ranges for each QSO where CIV can be detected.
    zmin, zmax, los, isHIRES = [], [], [], []
    for CIV in CIVa:
        isHIRES.append(CIV['filename'] in hires_filenames.values())
        zmax.append( (1 + CIV['zem']) * (1 - 3000. / Ckms) - 1 )
        zmin.append( wlya * (1 + CIV['zem']) / wc4 - 1 )
        los.append( cosmo.Dc(CIV['z']) * const )

    CIVb = rec_append_fields(CIVa, ['rew', 'zmin', 'zmax', 'dlos', 'isHIRES'],
                             [CIVa.ew / (1+CIVa.z), zmin, zmax, los, isHIRES])
    
    return  CIVb

def make_random_CIV(CIVall, nrand=None, trypickle=False):
    """ Generate random absorbers. For each QSO containing absorbers,
    generate the number of random absorbers as reals, but with a
    randomly selected redshift

    generate `nrand` times more randoms than reals."""

    random_CIV = []
    const = 1. / (PC * 1.e6)

    print 'Generating random CIV lines'
    # this loop takes around 10s to run for nrand=1000
    idnum = 1
    for CIV in CIVall:
        z0,z1 = CIV['zmin'], CIV['zmax']
        filename = CIV['filename']
        zran = z0 + (z1 - z0) * np.random.rand(nrand)
        row = [CIV[k] for k in 'name ra dec zem'.split()]
        rows = [tuple([i + idnum] + row + [z, cosmo.Dc(z) * const, filename])
                for i,z in enumerate(zran)]
        idnum += len(zran)
        random_CIV.extend(rows)

    random_CIV = np.rec.fromrecords(
        random_CIV, names='num,name,ra,dec,zem,z,dlos,filename')

    return random_CIV

def plotmatrix(mat, title=None, prefix=None, vmin=None, vmax=None):
    fig = pl.figure(figsize=A4LANDSCAPE)
    ax = fig.add_subplot(111)
    col = ax.pcolor(mat, vmin=vmin, vmax=vmax, cmap=pl.cm.Spectral_r)
    ax.set_xlabel('Transverse Distance / ($h^{-1}$ Mpc)')
    ax.set_ylabel('Radial Distance / ($h^{-1}$ Mpc)')
    ax.set_title(title)
    ax.axis('equal')
    ax.set_ylim(0,15)
    ax.set_xlim(0,15)
    fig.colorbar(col, fraction=0.22)
    if prefix is not None:
        fig.savefig('pdf/' + prefix + '.pdf')
        fig.savefig(prefix + '.png')

def plot1d(pairs, randompairs, nrandom=1000,
           cbins=None, title=None, prefix=None):
    """ Do plot of a 1-d xi (e.g. xi(s)) for a series of distance
    bins. Also show the number of real and random pairs in each bin.
    """
    fig = pl.figure(figsize=(8,10))
    fig.subplots_adjust(hspace=0.0001)
    ax = fig.add_subplot(211)#( (0.1, 0.1, 0.8, 0.8) )
    npairs = np.array([len(p) for p in pairs], float)
    err = np.sqrt(npairs)
    npairs_random = np.array([len(p) for p in randompairs], float)
    vals = npairs / npairs_random * nrandom - 1
    err = err / npairs_random * nrandom
    ax.plot(cbins, vals,'o', label=None)
    ax.errorbar(cbins, vals, yerr=err, fmt=None, ecolor='b',label=None)
    ax.set_ylabel(r'$\xi(s)$, real pairs / random pairs - 1')
    ax.set_title(title)
    ax.loglog(cbins, xifunc(cbins, 3.17, -1.60), '-r+',label='A03 fit')
    ax.set_xlim(0, 50)
    ax.legend()
    ax.xaxis.grid(which='minor')
    ax.grid()
    ax.xaxis.set_ticklabels([])

    ax = fig.add_subplot(212,sharex=ax)
    ax.loglog(cbins,npairs,ls='steps-mid',label='Real')
    ax.plot(cbins,npairs_random/nrandom,'r',ls='steps-mid',label='Random')
    ax.set_ylabel('Number of pairs')
    ax.set_xlabel('Comoving separation / ($h^{-1}$ Mpc)')
    ax.set_xlim(0, 50)
    ax.legend(loc='lower right')
    ax.xaxis.grid(which='minor')
    ax.grid()

    if prefix is not None:
        fig.savefig('pdf/' + prefix + '.pdf')
        fig.savefig(prefix + '.png')
    return vals, err

def xifunc(r, r0, gamma):
    """ Function often fitted to correaltion function, it's a
    prediction for the correaltion function in the linear clustering
    regime."""
    return (r / float(r0)) ** gamma


if 0:
    # read data
    lbgs = xcorr.readlbgs()
    CIVall0 = readCIV()
    CIVall = CIVall0[CIVall0.z < CIVall0.zmax]


if 0:
    #CIVa = CIVall[CIVall.isHIRES]
    #CIVa = CIVall[CIVall.rew < 0.3]
    CIVa = CIVall
    
    # generate random absorbers
    random_CIV = make_random_CIV(CIVa, nrand=NRAND, trypickle=False)
    
    # loop over every QSO to calculate number of pairs of CIV systems
    # and lbgs as a function of distance
    # need a unique specifier for each qso, filename acts as this
    
    qsos = list(set(CIVa.filename))

    # ids of CIV absorbers, galaxies that are pair members
#     pairIDs = xcorr_CIV.list_matrix(NBINS_R, NBINS_T)
#     paircounts = np.zeros([NBINS_R, NBINS_T], float)
#     random_pairIDs = xcorr_CIV.list_matrix(NBINS_R, NBINS_T)
#     random_paircounts = np.zeros([NBINS_R, NBINS_T], float)

#     pairIDs_t = [[] for i in range(NBINS_T)]
#     paircounts_t = np.zeros(NBINS_T, float)
#     random_pairIDs_t = [[] for i in range(NBINS_T)]
#     random_paircounts_t = np.zeros(NBINS_T, float)

    pairIDs_s = [[] for i in range(NBINS)]
    paircounts_s = np.zeros(NBINS, float)
    random_pairIDs_s = [[] for i in range(NBINS)]
    random_paircounts_s = np.zeros(NBINS, float)

    contributing_qsos = []
    for qso in qsos:
        # find all the CIV systems for this QSO
        c4 = CIVa[CIVa.filename == qso]
        if len(c4) == 0:
            continue
        zmin, zmax = c4.zmin[0], c4.zmax[0]
        print zmin, zmax
        ra, dec = c4.ra[0], c4.dec[0]

        c4ran = random_CIV[random_CIV.filename == qso]
        
        # get all the LBGs within
        dv = 5000.
        # go a little further beyond the z limits to look for LBGs.
        zmin = (1 + zmin) * (1 - dv/Ckms) - 1
        zmax = (1 + zmax) * (1 + dv/Ckms) - 1
        print 'Redshift limits for LBGs %5.2f -> %5.2f' % (zmin,zmax)
        
        # lbg selection
        c0 = (zmin <= lbgs['z']) & (lbgs['z'] <= zmax)
        lbgs0 = lbgs[c0]
        sep = ang_sep(ra, dec, lbgs0.ra, lbgs0.dec)
        is_close =  sep < MAXSEP
        lbgs1 = lbgs0[is_close]
        # if there are no LBGs close enough, skip to the next QSO.
        temp = qso.replace('/home/nhmc/projects/gas_gal_hiz/qsos/','')
        if len(lbgs1) == 0:
            print 'No LBGs close enough to %s, skipping' % temp
            continue
        else:
            contributing_qsos.append( (ra,dec,temp) )
        # angular separation of lbgs from QSO sightline (radians)
        lbg_angsep = sep[is_close] * np.pi / 180.
        # comoving transverse separation
        lbg_angdist = np.array([r*los for r,los in
                                zip(lbg_angsep, lbgs1.losdist)])

        ######################################
        # for xi(r,t)
        ######################################
        
#         # real pairs
#         pairs, pairids = xcorr_CIV.calc_npairs(
#             radial_binedges, transverse_binedges, lbg_angdist,
#             lbgs1.losdist, lbgs1.num, c4.dlos, c4.num)
#         paircounts += pairs
#         xcorr_CIV.list_matrix_extend(pairIDs, pairids)

#         # now random
#         pairs, pairids = xcorr_CIV.calc_npairs(
#             radial_binedges, transverse_binedges, lbg_angdist,
#             lbgs1.losdist, lbgs1.num, c4ran.dlos, c4ran.num)
#         random_paircounts += pairs
#         xcorr_CIV.list_matrix_extend(random_pairIDs, pairids)

        ######################################
        # for xi(s)
        ######################################
        pairs, pairids = calc_npairs_s(
            BINEDGES, lbg_angdist,
            lbgs1.losdist, lbgs1.num, c4.dlos, c4.num,
            maxdist)
        paircounts_s += pairs
        for j,p in enumerate(pairids):
            pairIDs_s[j].extend(p)
         
        pairs, pairids = calc_npairs_s(
            BINEDGES, lbg_angdist,
            lbgs1.losdist, lbgs1.num, c4ran.dlos, c4ran.num,
            maxdist)
        random_paircounts_s += pairs
        for j,p in enumerate(pairids):
            random_pairIDs_s[j].extend(p)

        ######################################
        # for xi(t)
        ######################################
#         pairs, pairids = calc_npairs_t(
#             transverse_binedges, lbg_angdist,
#             lbgs1.losdist, lbgs1.num, c4.dlos, c4.num)
#         paircounts_t += pairs
#         list_matrix_extend(pairIDs_t, pairids)
        
#         pairs, pairids = calc_npairs_t(
#             transverse_binedges, lbg_angdist,
#             lbgs1.losdist, lbgs1.num, c4ran.dlos, c4ran.num)
#         random_paircounts_t += pairs
#         list_matrix_extend(random_pairIDs_t, pairids)

if 0:
    #print contributing_qsos

    # look at xi(s) first...
    plot1d(pairIDs_s, random_pairIDs_s, cbins=CBINS,
           #title=r'CIV(EW < 0.3$\AA$)-LBG $\xi(s)$',
           title=r'CIV(all)-LBG $\xi(s)$',
           prefix='CIVlbg_xi_s_logbins')

    idgals = xcorr.unique_lbgids(pairIDs_s)
    idCIVs = xcorr.unique_absids(pairIDs_s)
    # ids of galaxies in the first, say, 5 bins:
    #idgals = unique_lbgids(pairIDs_s[:5])


if 0:
    # now find xi
    xi = NRAND * paircounts / random_paircounts - 1

    plotmatrix(paircounts, prefix='CIVpaircounts',
               title='Number of LBG - real CIV pairs')
    plotmatrix(random_paircounts / float(NRAND),
               prefix='CIVrandom_paircounts',
               title='Average number of LBG - random CIV pairs')
    plotmatrix(np.nan_to_num(xi), vmin=-1, prefix='CIVxi',
               title=r'$\xi(r_{\theta}$, $r_{z})$ (real / random - 1)')
    
    xi_t = NRAND * paircounts_t / random_paircounts_t - 1

    plot1d(pairIDs_s, random_pairIDs_s, title=r'CIV-LBG $\xi(s)$',
           prefix='CIVlbg_xi_s')
    show()
